import sys
parent_dir = "../pygl"
sys.path.append(parent_dir)
import ctypes
import time
import OpenGL.GL as gl
import glm
import imgui
import numpy as np
import pygl
from common import AppWithTrackball, Sphere, AxesHelper, GridHelper, InstancedSphere

from usecase.chem import ALLELEMENTS, Molecule, XYZReader

class SphereApp(AppWithTrackball):
    # 初始化着色器，vao vbo等
    def setupVertex(self) -> None:
        gl.glEnable(gl.GL_DEPTH_TEST)
        self.start_process = time.time()
        self.light_program = pygl.ProgramVF("intermediate/shaders/geometry.vs", 
                                      "intermediate/shaders/geometry.fs")
        self.program = pygl.ProgramVF("usecase/shaders/molecule.vs", 
                                      "usecase/shaders/molecule.fs")
        self.background_color = glm.vec4(0.5, 0.5, 0.5, 1.0)
        self.world_axes = AxesHelper(30, 2.0)
        self.grid = GridHelper(220.0, 30, glm.vec3(1.0, 1.0, 1.0), 
                               glm.vec3(0.2, 0.2, 0.2))
        self.show_world_axes = True
        self.show_grid = True
        self.imgui_window_width = 250
        self.imgui_window_height = 300
        self.molecule:Molecule = XYZReader.Open("dataset/4QCI.xyz")
        self.number_atom = len(self.molecule.atoms)
        atom_positions = np.empty((self.number_atom, 7), dtype=gl.GLfloat)
        index = 0
        for atom in self.molecule.atoms:
            color = self.RGBToGLColor(ALLELEMENTS[atom.number].CPK_color)
            atom_positions[index, 0] = atom.x
            atom_positions[index, 1] = atom.y
            atom_positions[index, 2] = atom.z
            atom_positions[index, 3] = color.x
            atom_positions[index, 4] = color.y
            atom_positions[index, 5] = color.z
            atom_positions[index, 6] = ALLELEMENTS[atom.number].vdw_radius
            index += 1

        self.atoms = InstancedSphere(atom_positions)
        self.light = Sphere()
        self.light_color = glm.vec3(1.0, 1.0, 1.0)
        self.light.position = glm.vec3(10.0, 100.0, 0.0)
        self.ambient_color = glm.vec3(1.0, 1.0, 1.0)
        self.Ka = 0.4
        self.Kd = 0.9
        self.Ks = 1.0
        self.shininess = 60.0
        self.camera.far = 1000
        self.camera.zoom_speed = 2.0
        self.camera.reset(glm.vec3(200.0, 260.0, 240.0), glm.vec3(0.0, 0.0, 0.0), 
                    glm.vec3(0.0, 1.0, 0.0))

    # 配置imgui界面元素
    def showImGui(self)->None:
        imgui.begin("设置")
        if imgui.button("重置轨迹球相机"):
            self.camera.reset(glm.vec3(200.0, 260.0, 240.0), glm.vec3(0.0, 0.0, 0.0), 
                    glm.vec3(0.0, 1.0, 0.0))
        _, self.show_world_axes = imgui.checkbox("显示三维世界坐标轴", 
                                                self.show_world_axes)
        _, self.show_grid = imgui.checkbox("显示XZ平面网格", 
                                                self.show_grid)
        imgui.end()       

    # 渲染图形
    def render(self) -> None:
        gl.glClearBufferfv(gl.GL_COLOR, 0, self.background_color.to_tuple())
        gl.glClearBufferfv(gl.GL_DEPTH, 0, 1.0)
        self.view = self.camera.getViewMatrix()
        self.projection = self.camera.getProjectionMatrix()

        # 代表灯位置的球体
        self.light_program.use()
        self.light.updateModelMatrix()
        self.light_program.setUniformMatrix4fv(0, self.light.model_matrix)
        self.light_program.setUniformMatrix4fv(1, self.view)
        self.light_program.setUniformMatrix4fv(2, self.projection)
        self.light_program.setUniform1i(3, 1)
        self.light_program.setUniform3fv(4, self.light_color)
        self.light.render()    
        # 绘制分子
        self.program.use()
        self.program.setUniform3fv(4, self.light_color)
        self.program.setUniform3fv(5, self.light.position)
        self.program.setUniform3fv(6, self.ambient_color)
        self.program.setUniform3fv(7, self.camera.position)
        self.program.setUniform1f(8, self.Ka)
        self.program.setUniform1f(9, self.Kd)
        self.program.setUniform1f(10, self.Ks)
        self.program.setUniform1f(11, self.shininess)
        self.atoms.updateModelMatrix()
        self.program.setUniformMatrix4fv(0, self.atoms.model_matrix)
        self.program.setUniformMatrix4fv(1, self.view)
        self.program.setUniformMatrix4fv(2, self.projection)
        self.atoms.render()

        if self.show_world_axes:
            self.world_axes.updateModelMatrix
            self.world_axes.render(self.view, self.projection)
        if self.show_grid:
            self.grid.render(self.view, self.projection)

    # 释放资源
    def cleanup(self) -> None:
        self.world_axes.delete()
        self.grid.delete()
        self.light.delete()
        self.atoms.delete()
        self.program.delete()
        self.light_program.delete()

    def RGBToGLColor(self, rgb:str) -> glm.vec3:
        red = int(rgb[1:3], base=16)
        green = int(rgb[3:5], base=16)
        blue = int(rgb[5:7], base=16)
        return glm.vec3(red/255.0, green/255.0, blue/255.0)

if __name__ == "__main__":
    app = SphereApp(100, 100, 1000, 800, "你好，蛋白质App！")
    app.run()


